#!/usr/bin/env node
/**
 * Electron Cookie提取工具
 * 用于从运行中的Electron应用程序提取Cookie
 */

const { program } = require('commander');
const chalk = require('chalk');
const fs = require('fs');
const path = require('path');
const os = require('os');
const psList = require('ps-list');
const findProcess = require('find-process');

// 导入自定义模块
const { findElectronProcesses } = require('./lib/process-finder');
const { extractCookiesFromProcess } = require('./lib/cookie-extractor');
const { findCookieFiles, readCookieFile } = require('./lib/file-scanner');
const { injectFridaScript } = require('./lib/memory-scanner');

// 创建输出目录
const OUTPUT_DIR = path.join(process.cwd(), 'cookies_extracted');
if (!fs.existsSync(OUTPUT_DIR)) {
  fs.mkdirSync(OUTPUT_DIR, { recursive: true });
}

// 设置命令行选项
program
  .name('electron-cookie-extractor')
  .description('从运行中的Electron应用程序提取Cookie')
  .version('1.0.0');

// 查找进程命令
program
  .command('find')
  .description('查找可能的Electron应用进程')
  .option('-n, --name <name>', '按名称过滤进程')
  .option('-d, --detailed', '显示详细信息')
  .action(async (options) => {
    console.log(chalk.blue('正在查找Electron应用进程...'));
    
    try {
      const processes = await findElectronProcesses(options.name);
      
      if (processes.length === 0) {
        console.log(chalk.yellow('未找到Electron应用进程'));
        return;
      }
      
      console.log(chalk.green(`找到 ${processes.length} 个Electron应用进程：`));
      
      for (const proc of processes) {
        console.log(chalk.white(`PID: ${proc.pid}, 名称: ${proc.name}`));
        
        if (options.detailed) {
          console.log(chalk.gray(`  命令行: ${proc.cmd}`));
          
          // 查找Cookie文件
          const cookieFiles = await findCookieFiles(proc);
          if (cookieFiles.length > 0) {
            console.log(chalk.cyan('  可能的Cookie文件:'));
            cookieFiles.forEach(file => console.log(`    ${file}`));
          }
          
          console.log(''); // 空行分隔
        }
      }
    } catch (err) {
      console.error(chalk.red(`发生错误: ${err.message}`));
    }
  });

// 提取Cookie命令
program
  .command('extract')
  .description('从指定进程提取Cookie')
  .argument('<pid>', '目标进程ID')
  .option('-m, --method <method>', '提取方法 (files, memory, all)', 'all')
  .option('-o, --output <dir>', '输出目录', OUTPUT_DIR)
  .action(async (pid, options) => {
    console.log(chalk.blue(`正在从进程 ${pid} 提取Cookie...`));
    
    try {
      // 确保输出目录存在
      if (!fs.existsSync(options.output)) {
        fs.mkdirSync(options.output, { recursive: true });
      }
      
      // 找到目标进程
      const processes = await findProcess('pid', parseInt(pid, 10));
      if (processes.length === 0) {
        console.log(chalk.red(`未找到PID为 ${pid} 的进程`));
        return;
      }
      
      const targetProcess = processes[0];
      console.log(chalk.green(`目标进程: ${targetProcess.name} (PID: ${targetProcess.pid})`));
      
      // 根据选择的方法提取Cookie
      let cookies = [];
      
      if (options.method === 'files' || options.method === 'all') {
        console.log(chalk.blue('从文件中提取Cookie...'));
        const cookieFiles = await findCookieFiles({ pid: targetProcess.pid, name: targetProcess.name });
        
        for (const file of cookieFiles) {
          console.log(chalk.cyan(`正在读取文件: ${file}`));
          try {
            const fileCookies = await readCookieFile(file);
            if (fileCookies.length > 0) {
              console.log(chalk.green(`从文件中提取了 ${fileCookies.length} 个Cookie`));
              cookies = cookies.concat(fileCookies);
            }
          } catch (err) {
            console.error(chalk.red(`无法读取Cookie文件 ${file}: ${err.message}`));
          }
        }
      }
      
      if (options.method === 'memory' || options.method === 'all') {
        console.log(chalk.blue('从内存中提取Cookie...'));
        try {
          const memoryCookies = await injectFridaScript(targetProcess.pid);
          if (memoryCookies.length > 0) {
            console.log(chalk.green(`从内存中提取了 ${memoryCookies.length} 个Cookie`));
            cookies = cookies.concat(memoryCookies);
          }
        } catch (err) {
          console.error(chalk.red(`从内存提取Cookie时出错: ${err.message}`));
        }
      }
      
      // 保存结果
      if (cookies.length > 0) {
        // 去重
        const uniqueCookies = [...new Map(cookies.map(c => [c.name + c.value, c])).values()];
        console.log(chalk.green(`共提取了 ${uniqueCookies.length} 个唯一Cookie`));
        
        // 保存为JSON
        const timestamp = new Date().toISOString().replace(/:/g, '-').replace(/\..+/, '');
        const outputFile = path.join(options.output, `cookies_${targetProcess.pid}_${timestamp}.json`);
        
        fs.writeFileSync(outputFile, JSON.stringify(uniqueCookies, null, 2));
        console.log(chalk.green(`Cookie已保存到: ${outputFile}`));
        
        // 创建分类文件
        const categories = {
          authentication: [],
          session: [],
          tracking: [],
          other: []
        };
        
        for (const cookie of uniqueCookies) {
          const name = cookie.name.toLowerCase();
          if (name.includes('token') || name.includes('auth') || name.includes('login')) {
            categories.authentication.push(cookie);
          } else if (name.includes('session') || name.includes('sid')) {
            categories.session.push(cookie);
          } else if (name.includes('track') || name.includes('visitor') || name.includes('uid')) {
            categories.tracking.push(cookie);
          } else {
            categories.other.push(cookie);
          }
        }
        
        // 保存分类结果
        const categoryFile = path.join(options.output, `categories_${targetProcess.pid}_${timestamp}.json`);
        fs.writeFileSync(categoryFile, JSON.stringify(categories, null, 2));
        console.log(chalk.green(`分类结果已保存到: ${categoryFile}`));
      } else {
        console.log(chalk.yellow('未提取到任何Cookie'));
      }
    } catch (err) {
      console.error(chalk.red(`发生错误: ${err.message}`));
    }
  });

// 简单的帮助命令
program
  .command('help')
  .description('显示帮助信息')
  .action(() => {
    console.log(chalk.cyan('Electron Cookie提取工具使用指南'));
    console.log(chalk.white('这个工具可以从运行中的Electron应用程序中提取Cookie，包括从磁盘和内存中获取。'));
    console.log('');
    console.log(chalk.white('常用命令:'));
    console.log(chalk.green('  node index.js find          ') + '查找Electron应用进程');
    console.log(chalk.green('  node index.js find -d       ') + '查找Electron应用进程并显示详细信息');
    console.log(chalk.green('  node index.js extract <pid> ') + '从指定进程提取Cookie');
    console.log('');
    console.log(chalk.white('使用步骤:'));
    console.log(chalk.white('1. 先使用find命令查找目标进程'));
    console.log(chalk.white('2. 记下目标进程的PID'));
    console.log(chalk.white('3. 使用extract命令提取Cookie'));
    console.log('');
    console.log(chalk.yellow('注意: 提取内存Cookie需要Frida支持，并可能需要管理员/root权限'));
  });

// 处理无命令情况
if (process.argv.length <= 2) {
  program.help();
}

// 解析命令行参数
program.parse(); 